
from Agents.MainGraph import graph
from langchain_core.messages import ToolMessage
from tools.utility import _print_event
import uuid
import traceback
# try:
#     graph.get_graph().draw_mermaid_png(output_file_path="main_graph.png")
# except Exception:
#     # This requires some extra dependencies and is optional
#     pass
thread_id = str(uuid.uuid4())


config = {
    "configurable": {
        "customer_id": 1,
        "thread_id": thread_id,
    }
}
_printed = set()
while True:
    try:
        user_input = input("请输入:")
        events = graph.stream(
            {"messages": ("user", user_input)}, config, stream_mode="values"
        )
        for event in events:
            _print_event(event, _printed)
        
        snapshot = graph.get_state(config)
        while snapshot.next:
            try:
                user_input = input(
                    "你赞成上述行动吗？输入‘是’继续；否则，输入‘否’。\n\n"
                )
            except:
                user_input = "是"

            if user_input.strip() == "是":
                # Just continue

                result = graph.invoke(
                    None,
                    config,
                )
                _print_event(result, _printed)
            else:
                result = graph.invoke(
                    {
                        "messages": [

                            ToolMessage(
                                tool_call_id=event["messages"][-1].tool_calls[0]["id"],
                                content=f"用户拒绝API调用,继续帮助用户，考虑用户的输入。",
                            )
                        ]
                    },

                    config,
                )
            snapshot = graph.get_state(config)

    except Exception as e:
        print(traceback.format_exc())